package cn.com.zetatech.maple.writer.common;

import com.zaxxer.hikari.HikariDataSource;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * @author cheng.cao@zetatech.com.cn
 */
@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class DataSourceHolder {
    private Map<String, DataSource> holder = new HashMap<>();
    private AtomicBoolean bindHook = new AtomicBoolean(false);
    private ReadWriteLock rwlock = new ReentrantReadWriteLock();

    public DataSource getDataSource(String url, String driver, String username, String password) {
        rwlock.readLock().lock();
        String key = url + "_" + driver;
        DataSource value = holder.get(key);
        try {
            if (value == null) {
                rwlock.readLock().unlock();
                rwlock.writeLock().lock();
                try {
                    if (null == (value = holder.get(key))) {

                        HikariDataSource ds = new HikariDataSource();
                        ds.setJdbcUrl(url);
                        ds.setDriverClassName(driver);
                        ds.setUsername(username);
                        ds.setPassword(password);
                        ds.setConnectionTestQuery("select 1");
                        ds.setMaxLifetime(60000);
                        ds.setConnectionTimeout(60000);
                        ds.setIdleTimeout(300000);
                        ds.setMaximumPoolSize(15);
                        holder.put(key, ds);
                        value = ds;

                        // 绑定钩子函数:关闭连接池
                        if (bindHook.compareAndSet(false, true)) {
                            Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {
                                @Override
                                public void run() {
                                    log.info("清理数据库连接池:连接池个数：{}", holder.size());
                                    for (DataSource source : holder.values()) {
                                        try {
                                            ((HikariDataSource) source).close();
                                        } catch (Exception e) {
                                        }
                                    }
                                }
                            }));
                        }
                    }
                } finally {
                    rwlock.writeLock().unlock();
                    rwlock.readLock().lock();
                }
            }
        } finally {
            rwlock.readLock().unlock();
        }
        return value;
    }

    private static class Holder {
        private static DataSourceHolder instance = new DataSourceHolder();
    }

    public static DataSourceHolder me() {
        return Holder.instance;
    }
}
